#include <stdio.h>

#define BLOCK_SIZE  16          // submatrix size
#define N           512       // matrix size is N*N

__global__ void matMult ( float * a, float * b, int n, float * c )
{
	int bx = blockIdx.x;        
	int by = blockIdx.y;

	int tx = threadIdx.x;       
	int ty = threadIdx.y;

	
	int aBegin = n*BLOCK_SIZE*by;
	int aEnd = aBegin + n - 1;
	
	int aStep = BLOCK_SIZE;
	
	int bBegin = BLOCK_SIZE * bx;
	
	int bStep = BLOCK_SIZE * n;
	float sum = 0.0f;           

	for ( int ia = aBegin, ib = bBegin; ia <= aEnd; ia += aStep, ib += bStep )
	{
		__shared__ float as [BLOCK_SIZE][BLOCK_SIZE];
		__shared__ float bs [BLOCK_SIZE][BLOCK_SIZE];

		as [ty][tx] = a [ia + n * ty + tx];
		bs [ty][tx] = b [ib + n * ty + tx];

		__syncthreads();    

		for ( int k = 0; k < BLOCK_SIZE; k++ )
			sum += as [ty][k] * bs [k][tx];

		__syncthreads();
	}

	int ic = n * BLOCK_SIZE * by + BLOCK_SIZE * bx;

	c [ic + n * ty + tx] = sum;
}

int main ( int argc, char *  argv [] )
{
	int numBytes = N * N * sizeof ( float );

	float * a = new float [N*N];
	float * b = new float [N*N];
	float * c = new float [N*N];

	for ( int i = 0; i < N; i++ )
		for ( int j = 0; j < N; j++ )
		{
			a [i] = 0.0f;
			b [i] = 1.0f;
		}

		float * adev = NULL;
		float * bdev = NULL;
		float * cdev = NULL;

		cudaMalloc ( (void**)&adev, numBytes );
		cudaMalloc ( (void**)&bdev, numBytes );
		cudaMalloc ( (void**)&cdev, numBytes );

		dim3 threads ( BLOCK_SIZE, BLOCK_SIZE );
		dim3 blocks  ( N / threads.x, N / threads.y);

		cudaEvent_t start, stop;
		float gpuTime = 0.0f;

		cudaEventCreate ( &start );
		cudaEventCreate ( &stop );

		cudaEventRecord ( start, 0 );
		cudaMemcpy      ( adev, a, numBytes, cudaMemcpyHostToDevice );
		cudaMemcpy      ( bdev, b, numBytes, cudaMemcpyHostToDevice );

		matMult<<<blocks, threads>>> ( adev, bdev, N, cdev );

		cudaMemcpy      ( c, cdev, numBytes, cudaMemcpyDeviceToHost );
		cudaEventRecord ( stop, 0 );

		cudaEventSynchronize ( stop );
		cudaEventElapsedTime ( &gpuTime, start, stop );

		printf("Time elapsed: %.2f millseconds\n", gpuTime );

		cudaEventDestroy ( start );
		cudaEventDestroy ( stop  );
		cudaFree         ( adev  );
		cudaFree         ( bdev  );
		cudaFree         ( cdev  );

		delete a;
		delete b;
		delete c;
		system("pause");
		return 0;
}
